package com.jiang.daily.Q3372;

import java.util.ArrayList;
import java.util.List;

/**
 * @author Jiang
 * @version 1.0
 * @date 2025/05/28 9:58
 */
class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int n = edges1.length + 1;
        int[] cnt1 = cnt(edges1, k);
        int[] cnt2 = cnt(edges2, k - 1);

        int maxCnt2 = 0;
        for (int c: cnt2) {
            maxCnt2 = Math.max(maxCnt2, c);
        }
        int[] res = new int[n];
        for (int i = 0; i < n; i++) {
            res[i] = cnt1[i] + maxCnt2;
        }
        return res;
    }

    /**
     * 统计每个节点满足的子节点数
     * @param edges 边
     * @param k 层数
     * @return 每个节点满足的子节点数
     */
    private int[] cnt(int[][] edges, int k) {
        int len = edges.length + 1;
        List<List<Integer>> list = new ArrayList<>();
        for (int i = 0; i < len; i++) {
            list.add(new ArrayList<>());
        }

        // 每个节点的相邻节点
        for (int[] edge : edges) {
            list.get(edge[0]).add(edge[1]);
            list.get(edge[1]).add(edge[0]);
        }

        int[] res = new int[len];
        for (int i = 0; i < len; i++) {
            res[i] = dfs(i, -1, k, list);
        }
        return res;
    }

    /**
     * 递归求解
     * @param cur 当前节点
     * @param pre 上一个节点
     * @param k 层数
     * @param list 节点的相邻节点
     * @return 当前节点满足的子节点数
     */
    private int dfs(int cur, int pre, int k, List<List<Integer>> list) {
        if (k < 0) {
            return 0;
        }
        if (k == 0) {
            return 1;
        }
        int res = 1;
        for (int next : list.get(cur)) {
            if (next == pre) {
                continue;
            }
            res += dfs(next, cur, k - 1, list);
        }
        return res;
    }
}
